import requests
import json
import os


def send_chat_message(api_key, query, conversation_id="", user_id="abc-123", files=None, stream=False):
    """
    发送聊天消息到API端点

    :param api_key: API认证密钥
    :param query: 用户查询内容
    :param conversation_id: 可选，会话ID
    :param user_id: 可选，用户ID
    :param files: 可选，文件列表
    :param stream: 是否使用流式响应
    :return: 流式模式下返回生成器，非流式模式下返回响应对象
    """
    url = "http://8.136.111.93:8008/v1/chat-messages"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    # 从JSON文件读取数据
    json_file = os.path.join(os.path.dirname(
        __file__), 'fill_in_blank_question.json')
    with open(json_file, 'r', encoding='utf-8') as f:
        question_data = json.load(f)

    payload = {
        "inputs": {
            "courseName": "生物",
        },
        "query": json.dumps(question_data, ensure_ascii=False),
        "response_mode": "streaming" if stream else "blocking",
        "conversation_id": conversation_id,
        "user": user_id,
        "files": files or []
    }

    if stream:
        def generate():
            full_response = []
            with requests.post(url, headers=headers, json=payload, stream=True) as response:
                for chunk in response.iter_content(chunk_size=None):
                    if chunk:
                        decoded_chunk = chunk.decode('utf-8')
                        full_response.append(decoded_chunk)
                        yield decoded_chunk
            # 返回完整数据作为生成器的最后一个元素
            yield "[完整数据开始]\n" + ''.join(full_response) + "\n[完整数据结束]"

        return generate()
    else:
        response = requests.post(url, headers=headers, json=payload)
        return response


if __name__ == "__main__":
    api_key = "app-Zi524Vvwmtfn9xKwwtJ9r9ul"
    query = "智能补全"

    print("\nstarttttt")
    # 流式模式使用
    generator = send_chat_message(api_key, query, stream=True)
    full_data = None

    for chunk in generator:
        print(chunk, end='', flush=True)
        full_data = chunk  # 最后一个chunk就是完整数据

    # 添加分隔线便于识别
    print("\n\n======= 完整数据 =======")
    print(full_data)
    print("=======================")

    # 将完整数据保存到文件
    output_file = os.path.join(os.path.dirname(__file__), 'response_data.txt')
    with open(output_file, 'w', encoding='utf-8') as f:
        # 确保 full_data 是字符串类型
        if isinstance(full_data, bytes):
            full_data = full_data.decode('utf-8')
        elif full_data is None:
            full_data = ""
        f.write(full_data)
    print(f"\n数据已保存到: {output_file}")
